#include <cstdlib>
#include <iostream>
#include <cstdio>
#include <algorithm>
#include <climits>
#include "opencv2/myopencv.hpp"
#include "augmentopencv.h"
#include "SimpleHistogram.h"
#include "DepthTracker.h"

using namespace std;

string depthName, videoName;

// double lb = 6000, ub = 7500;
double lb = 0, ub = 32767;
double rg=ub-lb;

int videoCurFrame;

DepthTracker *tracker;

class TRectInfo {
public:
	CvPoint lu;
	CvPoint rb;
	double depthValue;
	CvRect rect;
	bool pressed;
	CvMat *canvas;
	string name;

	TRectInfo(void) {
		pressed = false;
		canvas = NULL;
		depthValue = 0;
	}
	~TRectInfo(void) {
		if (canvas != NULL)
			cvReleaseMat(&canvas);
	}
};
TRectInfo videoRect, depthRect;
class DepthPlayer* curPlayer = NULL;

class DepthPlayer{
public:
	FILE* fin;
	long long startFilePos;
	long long fileSize;
	int w, h;
	int frameNum;
	int curFrame;
	int frameSize;
	int seconds;
	short* frame;
	CvMat* mat;
	CvMat* newmat;
	CvMat* backup;
	bool isSkipTo;
	int restart;
	bool pause;
	DepthPlayer() {
		fin = NULL;
		startFilePos = fileSize = 0;
		w = h = 0;
		frameNum = curFrame = frameSize = 0;
		seconds = 0;
		frame = NULL;
		mat = newmat = backup = NULL;
		isSkipTo = true;
		restart = 0;
		pause = false;
	}

	void initParamter(string name) {
		fopen_s(&fin, name.c_str(), "rb");
		fscanf_s(fin, "%d %d\n", &w, &h);
		startFilePos = _ftelli64(fin);
		_fseeki64(fin, 0, SEEK_END);
		fileSize = _ftelli64(fin);
		_fseeki64(fin, startFilePos, SEEK_SET);

		frameSize = w*h*sizeof(short);
		frameNum = int((fileSize-startFilePos) / frameSize);
		seconds = frameNum / 30;

		frame = new short[w*h];
		mat = cvCreateMatHeader(h, w, CV_16SC1);
		newmat = cvCreateMat(h, w, CV_64FC1);
		backup = cvCreateMat(h, w, CV_16SC1);
		cvZero(backup);
	}

	void myflip( const uchar* src, size_t sstep, uchar* dst, size_t dstep, cv::Size size, size_t esz )
	{
		int i, j, limit = (int)(((size.width + 1)/2)*esz);
		cv::AutoBuffer<int> _tab(size.width*esz);
		int* tab = _tab;

		for( i = 0; i < size.width; i++ )
			for( size_t k = 0; k < esz; k++ )
				tab[i*esz + k] = (int)((size.width - i - 1)*esz + k);

		for( ; size.height--; src += sstep, dst += dstep )
		{
			for( i = 0; i < limit; i++ )
			{
				j = tab[i];
				uchar t0 = src[i], t1 = src[j];
				dst[i] = t1; dst[j] = t0;
			}
		}
	}

	void updateMat(void) {
		mat->data.s = frame;
		cv::Mat src = cv::cvarrToMat(mat);
		size_t esz = src.elemSize();
		myflip( src.data, src.step, src.data, src.step, src.size(), esz );
// 		cvSmooth(mat, mat);
// 		cvFlip(mat, NULL, 1);
	}

	void recovery(void)
	{
		short* matPtr;
		short* backupPtr;
		matPtr = mat->data.s;
		backupPtr = backup->data.s;
		for (int i=0; i<mat->rows; ++i)
			for (int j=0; j<mat->cols; ++j, ++matPtr, ++backupPtr) {
				if (*matPtr < 0) *matPtr = *backupPtr;
				if (*matPtr<lb || *matPtr>ub) *matPtr = 0;
			}
	}

	~DepthPlayer() {
		fclose(fin);
		delete[] frame;
		cvReleaseMat(&mat);
		cvReleaseMat(&newmat);
		cvReleaseMat(&backup);
	}
};

void showFrame(bool useUpdate = true) {

	if (curPlayer->pause)
		cout << "video frame number: " << videoCurFrame << "   "
			<< "depth frame number: " << curPlayer->curFrame << endl;

	if (useUpdate)
		curPlayer->updateMat();
	cvShowImage("depth tracking", curPlayer->mat);

	// 	cvConvertScale(mat, newmat, 1, -lb);
	// 	cvConvertScale(newmat, newmat, 1/rg, 0);
	// 	cvShowImage("original filtered data", newmat);

	// 	recovery();
	// 	kalmanRecovery();
	// 	cvConvertScale(mat, newmat, 1, -lb);
	// 	cvConvertScale(newmat, newmat, 1/rg, 0);
	// 	cvShowImage("recovery filtered data", newmat);
	// 	cvCopy(mat, backup);
	cvWaitKey(1000/60);
}

void skipTo(int frameNo) {
	if (curPlayer->isSkipTo == false) return;
	curPlayer->curFrame = frameNo;
	long long tmp = curPlayer->curFrame;
	tmp *= curPlayer->frameSize;
	tmp += curPlayer->startFilePos;
	_fseeki64(curPlayer->fin, tmp, SEEK_SET);
	if (curPlayer->pause) {
		fread_s(curPlayer->frame, curPlayer->frameSize, sizeof(short), 
			curPlayer->w*curPlayer->h, curPlayer->fin);
		showFrame();
	}
}

void skipToCallback(int val) {
	skipTo(val*30);
}

void restartCallback(int state, void* userdata) {
	curPlayer->restart = 1;
}

void pauseCallback(int state, void* userdata) {
	curPlayer->pause = !curPlayer->pause;
}

void prevFrameCallback(int state, void* userdata) {
	skipTo(curPlayer->curFrame-1);
}

void nextFrameCallback(int state, void* userdata) {
	skipTo(curPlayer->curFrame+1);
}

void prev10FrameCallback(int state, void* userdata) {
	skipTo(curPlayer->curFrame-10);
}

void next10FrameCallback(int state, void* userdata) {
	skipTo(curPlayer->curFrame+10);
}

int getDepthFrameNum(string name) {
	FILE* fin;
	int w, h;
	fopen_s(&fin, name.c_str(), "rb");
	fscanf_s(fin, "%d %d\n", &w, &h);
	long long startFilePos = _ftelli64(fin);
	_fseeki64(fin, 0, SEEK_END);
	long long fileSize = _ftelli64(fin);
	fclose(fin);

	int frameSize = w*h*sizeof(short);
	return int((fileSize-startFilePos) / frameSize);
}

int getVideoFrameNum(string name) {
	cv::VideoCapture video;
	video.open(name);
	if (video.isOpened() == false) return -1;
	return cvRound(video.get(CV_CAP_PROP_FRAME_COUNT));
}




void drawRect(int type, int x, int y, int flag, void* param) {
	TRectInfo *rectInfo = (TRectInfo*)param;
	CvPoint first, second;
	first = rectInfo->lu;
	second = rectInfo->rb;

	switch (type) {
	case CV_EVENT_RBUTTONDOWN:
		CvRect rect;
		rect.x = first.x;
		rect.width = second.x - first.x;
		rect.y = first.y;
		rect.height = second.y - first.y;
		tracker->showHist(curPlayer->backup, rect);
		break;
	case CV_EVENT_LBUTTONDOWN:
		first = cvPoint(x, y);
		rectInfo->pressed = true;
		break;
	case CV_EVENT_LBUTTONUP:
		{
			second = cvPoint(x, first.y+x-first.x);
			CvMat *canvas = cvCloneMat(rectInfo->canvas);
			CvFont font;
			cvInitFont(&font, CV_FONT_HERSHEY_SIMPLEX, 0.7, 0.7);
			char words[100];
			sprintf_s(words, 100, "first=(%d, %d), second=(%d, %d)", first.x, first.y, second.x, second.y);
			int maxInt = (1 << CV_ELEM_SIZE(canvas->type)*8) - 1;
			cvPutText(canvas, words, cvPoint(50,50), &font, cvScalar(maxInt));
			cvRectangle(canvas, first, second, cvScalar(maxInt));
			cvShowImage(rectInfo->name.c_str(), canvas);
			cvReleaseMat(&canvas);
		}
		break;
	case CV_EVENT_MOUSEMOVE:
		if (flag & CV_EVENT_FLAG_LBUTTON) {
			second = cvPoint(x, first.y+x-first.x);
			CvMat *canvas = cvCloneMat(rectInfo->canvas);
			CvFont font;
			cvInitFont(&font, CV_FONT_HERSHEY_SIMPLEX, 0.7, 0.7);
			char words[100];
			sprintf_s(words, 100, "first=(%d, %d), second=(%d, %d)", first.x, first.y, second.x, second.y);
			int maxInt = (1 << CV_ELEM_SIZE(canvas->type)*8) - 1;
			cvPutText(canvas, words, cvPoint(50,50), &font, cvScalar(maxInt));
			cvRectangle(canvas, first, second, cvScalar(maxInt));
			cvShowImage(rectInfo->name.c_str(), canvas);
			cvReleaseMat(&canvas);
		}
		break;
	default:
		break;
	}
	rectInfo->lu = first;
	rectInfo->rb = second;
}

void showFiltered(void) {
	CvMat *newmat = cvCreateMat(curPlayer->mat->rows, curPlayer->mat->cols, CV_32F);
	cvConvertScale(curPlayer->mat, newmat, 1, -lb);
	cvConvertScale(newmat, newmat, 1/rg, 0);
	cvShowImage("original filtered data", newmat);
	cvReleaseMat(&newmat);
}

inline int cvFitIn(int x, int a, int b) {return min(max(a,x),b);}

void bothPlay(string depthName, string videoName) {
	curPlayer = new DepthPlayer;
	curPlayer->initParamter(depthName);
	int val;
	cvNamedWindow("original depth data");
	cvCreateTrackbar("trackbar", "original depth data", &val, curPlayer->seconds, skipToCallback);
	cvCreateButton("restart", restartCallback, NULL, CV_PUSH_BUTTON);
	cvCreateButton("pause", pauseCallback, NULL, CV_CHECKBOX);
	cvCreateButton("previousFrame", prevFrameCallback, NULL, CV_PUSH_BUTTON);
	cvCreateButton("nextFrame", nextFrameCallback, NULL, CV_PUSH_BUTTON);
	cvCreateButton("prev10Frame", prev10FrameCallback, NULL, CV_PUSH_BUTTON);
	cvCreateButton("next10Frame", next10FrameCallback, NULL, CV_PUSH_BUTTON);
	cvResizeWindow("original depth data", 1, 1);
	cvMoveWindow("original depth data", -1000, 0);
	cvNamedWindow("original video data");
	cvMoveWindow("original video data", 20, 0);
	cvNamedWindow("depth tracking");
	cvMoveWindow("depth tracking", 700, 0);


	cv::VideoCapture video;
	video.open(videoName);

	curPlayer->curFrame = 0;
	int readFrames = 0;
	long long startTime;
	double cpufreq = cvGetTickFrequency();

	cv::Mat videoMat;
	cv::Mat videoMat16;
	cv::Mat backupVideoMat = cv::Mat::zeros(curPlayer->h, curPlayer->w, CV_16SC1);
	

	int videoStartFrame = 35;
	for (int i=0; i<videoStartFrame; ++i)
		video.read(videoMat);

	CvMat *depthRecovMat = cvCloneMat(curPlayer->backup);
	cv::Mat videoOriginMat16(curPlayer->h, curPlayer->w, CV_16SC1);
	

	bool firstShow = true;
	while (true) {
		if (readFrames % 10 == 0) {
			startTime = cvGetTickCount();	
		}
		if (curPlayer->curFrame % 30 == 0) {
			curPlayer->isSkipTo = false;
			cvSetTrackbarPos("trackbar", "original depth data", curPlayer->curFrame/30);
			curPlayer->isSkipTo = true;
		}
		fread_s(curPlayer->frame, curPlayer->frameSize, sizeof(short), 
			curPlayer->w*curPlayer->h, curPlayer->fin);

		video.read(videoMat);
		cv::cvtColor(videoMat, videoMat, cv::COLOR_RGB2GRAY);
		videoMat.convertTo(videoMat16, CV_16SC1);

		++readFrames;
		if (readFrames % 10 == 0) {
			cout << "FPS: " << 10 / ((cvGetTickCount() - startTime) / cpufreq / 1000000) << endl;
			video.read(videoMat);
			cv::cvtColor(videoMat, videoMat, cv::COLOR_RGB2GRAY);
			videoMat.convertTo(videoMat16, CV_16SC1);
		}
		if (readFrames % 100 == 0) {
			video.read(videoMat);
			cv::cvtColor(videoMat, videoMat, cv::COLOR_RGB2GRAY);
			videoMat.convertTo(videoMat16, CV_16SC1);
		}
		videoCurFrame = curPlayer->curFrame;
		while (curPlayer->pause) {
			cvWaitKey(1000/30);
		}

		curPlayer->updateMat();

// 		cvShowImage("original depth data", curPlayer->mat);

		curPlayer->recovery();

		cvCopy(curPlayer->mat, depthRecovMat);

		videoMat16.copyTo(videoOriginMat16);

		cvSub(curPlayer->mat, curPlayer->backup, curPlayer->mat);
		videoMat16 = cv::abs(videoMat16 - backupVideoMat);

		cvCopy(depthRecovMat, curPlayer->backup);
		videoOriginMat16.copyTo(backupVideoMat);

		if (!firstShow) {
			double depthVal = 0;
			CvRect rect = tracker->tracking(curPlayer->backup, &depthVal);
			lb = tracker->lb;
			ub = tracker->ub;
			{
				CvPoint first, second;
				first.x = rect.x;
				first.y = rect.y;
				second.x = first.x + rect.width;
				second.y = first.y + rect.height;
				cvCopy(curPlayer->backup, curPlayer->mat);
				cvRectangle(curPlayer->mat, first, second, cvScalar(32767));
				CvPoint3D64f offset;
				offset.x = depthRect.rect.x - rect.x;
				offset.y = depthRect.rect.y - rect.y;
				offset.z = depthRect.depthValue - depthVal;
				depthRect.rect = rect;
				depthRect.depthValue = depthVal;
				double depthPS = 5./3;
				double videoPS = 1/5.;
				offset.x = 1./2*offset.x*depthPS + sqrt(3.)/2*(videoPS*offset.z);

				videoRect.rect.x -= offset.x;
				videoRect.rect.y -= offset.y;
			}
			{
				cv::Point first, second;
				first.x = cvFitIn(videoRect.rect.x, 0, videoMat16.cols);
				first.y = cvFitIn(videoRect.rect.y, 0, videoMat16.rows);
				second.x = cvFitIn(first.x + videoRect.rect.width, 0, videoMat16.cols);
				second.y = cvFitIn(first.y + videoRect.rect.height, 0, videoMat16.rows);
				cv::rectangle(videoMat, first, second, cv::Scalar(255));
			}
		};

		showFrame(false);
		videoMat16.convertTo(videoMat16, CV_8UC1);
// 		imshow("video data", videoMat16);

		imshow("original video data", videoMat);

		if (firstShow) {
			char c=27;
			firstShow = false;
			cvNamedWindow("first video");
			cvSetMouseCallback("first video", drawRect, &videoRect);
			cvNamedWindow("first depth");
			cvSetMouseCallback("first depth", drawRect, &depthRect);

			while (c==27 || !depthRect.pressed || !videoRect.pressed) {
				if (!depthRect.pressed)
					cvShowImage("first depth", curPlayer->backup);
				if (!videoRect.pressed)
					imshow("first video", videoMat);

				depthRect.canvas = cvCloneMat(curPlayer->backup);
				CvMat tmp = videoMat;
				videoRect.canvas = cvCloneMat(&tmp);

				videoRect.name = "first video";
				depthRect.name = "first depth";

// 				c = cvWaitKey(0);
				videoRect.lu.x = 167;
				videoRect.lu.y = 173;
				videoRect.rb.x = 449;
				videoRect.rb.y = 455;
				depthRect.lu.x = 284;
				depthRect.lu.y = 216;
				depthRect.rb.x = 422;
				depthRect.rb.y = 354;
				depthRect.pressed = true;
				videoRect.pressed = true;
				c=0;

				cout << "char: " << (int)c << endl;
			}
			cvDestroyWindow("first depth");
			cvDestroyWindow("first video");
			CvRect rect;
			rect.x = videoRect.lu.x;
			rect.y = videoRect.lu.y;
			rect.width = videoRect.rb.x - rect.x;
			rect.height = videoRect.rb.y - rect.y;
			videoRect.rect = rect;

			rect.x = depthRect.lu.x;
			rect.y = depthRect.lu.y;
			rect.width = depthRect.rb.x - rect.x;
			rect.height = depthRect.rb.y - rect.y;
			depthRect.rect = rect;

			tracker->initTracking(curPlayer->backup, depthRect.rect, &depthRect.depthValue);
			lb = tracker->lb;
			ub = tracker->ub;

		}

		++curPlayer->curFrame;
		if (curPlayer->restart) {
			curPlayer->restart = 0;
			video.release();
			video.open(videoName);
			for (int i=0; i<videoStartFrame; ++i)
				video.read(videoMat);
			skipTo(0);
		}

		if (feof(curPlayer->fin)) {
			skipTo(0);
		}

		int c = cvWaitKey(10);
// 		cout << c << endl;
		if (c == 'd') {
			cvNamedWindow("debug depth");
			cvSetMouseCallback("debug depth", drawRect, &depthRect);
			cvShowImage("debug depth", curPlayer->backup);

			while (c == 'd') {
				if (depthRect.canvas != NULL)
					cvReleaseMat(&depthRect.canvas);
				depthRect.canvas = cvCloneMat(curPlayer->backup);
				depthRect.name = "debug depth";

				c = cvWaitKey(0);
				cout << "char: " << (int)c << endl;
			}
			cvDestroyWindow("debug depth");

		}

	}
	cvReleaseMat(&depthRecovMat);
}

int main(int argc, char* argv[])
{
	tracker = new DepthTracker;
	if (argc != 3) {
		cout << "need two arguments." << endl;
		return 1;
	}

	depthName = string(argv[1]);
	videoName = string(argv[2]);
	if (depthName[depthName.length()-1] == 'i')
		swap(depthName, videoName);

	cout << "depth file: " << depthName << endl << "video file: " << videoName << endl;

	cout << "depth frames: " << getDepthFrameNum(depthName) << endl
		<< "video frames: " << getVideoFrameNum(videoName) << endl;

	bothPlay(depthName, videoName);

	return 0;
}
